import torch


def masked_mse_loss(pred, target, mask):
    if (~mask).sum() == 0:
        return (pred * 0.).sum()
    masked_diff = torch.where(mask, torch.tensor(0.0, device=pred.device), pred - target)

    squared_error = masked_diff.pow(2)

    loss = squared_error.sum() / (~mask).sum()
    return loss


def masked_mae_loss(pred, target, mask):
    if (~mask).sum() == 0:
        return (pred * 0.).sum()

    masked_diff = torch.where(mask, torch.tensor(0.0, device=pred.device), pred - target)
    abs_error = masked_diff.abs()

    loss = abs_error.sum() / (~mask).sum()
    return loss